from pydantic import BaseModel
from typing import List, Optional

class RagQueryRequest(BaseModel):
    """RAG查询请求模型"""
    knowledge_bases: List[str]
    question: str

class ThinkingMessage(BaseModel):
    """思考中消息模型"""
    type: str = "thinking"
    text: str = "正在思考中，请稍候..."

class TextChunkMessage(BaseModel):
    """文本块消息模型"""
    type: str = "text"
    text: str

class ImageInfoMessage(BaseModel):
    """图片信息消息模型"""
    type: str = "images_info"
    images: List[dict]

class ErrorMessage(BaseModel):
    """错误消息模型"""
    type: str = "error"
    message: str

class EndMessage(BaseModel):
    """结束消息模型"""
    type: str = "end"